import akshare as ak
import stock_zt as zt
from datetime import datetime
import pandas as pd


# 优化数据
def optimize_data(df):
    """
    :param df:
    :return:
    """
    # 主力净流入-净额单位转为万元，并且保留2位小数
    df['主力净流入-净额'] = df['主力净流入-净额'] / 10000
    df['主力净流入-净额（万元）'] = df['主力净流入-净额'].round(2)
    del df['主力净流入-净额']

    # 超大单净流入-净额单位转换为万元，并且保留2位小数
    df['超大单净流入-净额'] = df['超大单净流入-净额'] / 10000
    df['超大单净流入-净额（万元）'] = df['超大单净流入-净额'].round(2)
    del df['超大单净流入-净额']

    # 大单净流入-净额单位转换为万元，并且保留2位小数
    df['大单净流入-净额'] = df['大单净流入-净额'] / 10000
    df['大单净流入-净额（万元）'] = df['大单净流入-净额'].round(2)
    del df['大单净流入-净额']

    # 中单净流入-净额单位转换为万元，并且保留2位小数
    df['中单净流入-净额'] = df['中单净流入-净额'] / 10000
    df['中单净流入-净额（万元）'] = df['中单净流入-净额'].round(2)
    del df['中单净流入-净额']

    # 小单净流入-净额单位转换为万元，并且保留2位小数
    df['小单净流入-净额'] = df['小单净流入-净额'] / 10000
    df['小单净流入-净额（万元）'] = df['小单净流入-净额'].round(2)
    del df['小单净流入-净额']

    return df


def get_zl_hold(code, name, market, day):
    # 获取股票资金流向数据
    all_data = ak.stock_individual_fund_flow(code, market)
    # 获取最近5天的数据
    get_past_5_days_data = pd.DataFrame(all_data.tail(day))
    # 优化数据
    get_past_5_days_data = optimize_data(get_past_5_days_data)
    # 添加代码列
    get_past_5_days_data['代码'] = code
    # 添加名称列
    get_past_5_days_data['名称'] = name
    return pd.DataFrame(get_past_5_days_data)


# 获取最近5天一个涨停的股票
data_past_5_days_first_zt = zt.data_past_5_days_first_zt
# 所有近五天数据
all_past_5_day_data = pd.DataFrame()
for i in data_past_5_days_first_zt.index:
    code = data_past_5_days_first_zt.loc[i, '代码']
    name = data_past_5_days_first_zt.loc[i, '名称']
    market = "sh" if code.startswith("6") else "sz"
    day = int(data_past_5_days_first_zt.loc[i, '距今涨停交易日']) + 1
    all_past_5_day_data = pd.concat([all_past_5_day_data, get_zl_hold(code, name, market,day)])
# 储存地址
path_data_zl_hold = zt.get_file_path('首板股票主力资金流动数据.csv')
# 储存数据
all_past_5_day_data.to_csv(path_data_zl_hold)

